{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "447b9876",
   "metadata": {},
   "source": [
    "This is part 1 of a tutorial series. We recommend following them in order, starting with [Part 0: Welcome to `musica`](0.%20Welcome%20to%20MUSICA.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4",
   "metadata": {},
   "source": [
    "# TS1 Box Model\n",
    "\n",
    "This tutorial walks through a coupled `tuv-x` / `micm` box model simulation using the TS1 (Tropospheric Standard 1) atmospheric chemistry mechanism.\n",
    "\n",
    "The simulation runs over a vertical column of 9 altitude grid cells (1–10 km) for a real geographic location (Boulder, CO) beginning at 6:30 AM local time. At each time step, TUV-x recomputes photolysis rate constants based on the current solar zenith angle, `ussa1976` provides temperature and pressure from the US Standard Atmosphere, and `micm` integrates the chemical kinetics forward in time.\n",
    "\n",
    "Before working through this tutorial, it is recommended to complete:\n",
    "- [Tutorial 1: Multiple Grid Cells](1.%20multiple_grid_cells.ipynb) — covers MICM state setup and multi-cell solving\n",
    "- [Tutorial 8: TUV-x Standard Configurations](8.%20tuv-x_standard_configurations.ipynb) — covers computing photolysis rate constants with TUV-x\n",
    "- [Tutorial 10: Chapman Cycle](10.%20chapman.ipynb) — covers coupling TUV-x and MICM with diurnal photolysis updates"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2c3d4e5",
   "metadata": {},
   "source": [
    "## 1. Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3d4e5f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from datetime import datetime, time, timedelta\n",
    "from zoneinfo import ZoneInfo\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pvlib\n",
    "import ussa1976\n",
    "import xarray as xr\n",
    "\n",
    "import musica\n",
    "from musica.mechanism_configuration import Parser\n",
    "from musica.micm.solver_result import SolverState\n",
    "from musica.tuvx import vTS1\n",
    "from musica.utils import find_config_path\n",
    "\n",
    "SECONDS_PER_HOUR = 3600"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c364473",
   "metadata": {},
   "source": [
    "## 2. Location and Simulation Time\n",
    "\n",
    "As in the [Chapman tutorial](10.%20chapman.ipynb), we simulate over Boulder, CO starting at 6:30 AM local time. All times are tracked as timezone-aware `datetime` objects in UTC."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c64640cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "boulder = (40.01879858223568, -105.27492413846649)  # (latitude, longitude)\n",
    "boulder_tz = ZoneInfo(\"America/Denver\")\n",
    "\n",
    "today_local = datetime.now(boulder_tz).date()\n",
    "sim_local  = datetime.combine(today_local, time(6, 30), tzinfo=boulder_tz)\n",
    "sim_time    = sim_local.astimezone(ZoneInfo(\"UTC\"))\n",
    "\n",
    "print(f\"Simulation start (UTC):   {sim_time}\")\n",
    "print(f\"Simulation start (local): {sim_time.astimezone(boulder_tz)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4e5f6a7",
   "metadata": {},
   "source": [
    "## 3. Load `tuv-x` and Define a Helper to Compute Photolysis Rates\n",
    "\n",
    "We load the vTS1 `tuv-x` calculator once and keep it alive for the duration of the simulation. The alias table in the TS1 config file maps TUV-x reaction labels to the `USER.j*` parameter names expected by the MICM mechanism.\n",
    "\n",
    "We define a helper function that:\n",
    "1. Uses `pvlib` to compute the solar zenith angle (SZA) at Boulder for a given time.\n",
    "2. Runs TUV-x with that SZA to get photolysis rate constants at every altitude.\n",
    "3. Maps the TUV-x labels to `USER.j*` names using the alias table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f6a7b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "tuvx = vTS1.get_tuvx_calculator()\n",
    "\n",
    "tuv_path = find_config_path(\"tuvx\", \"ts1_tsmlt.json\")\n",
    "with open(tuv_path, 'r') as f:\n",
    "    data = json.load(f)\n",
    "alias_mappings = data.get('__CAM options', {}).get('aliasing', {}).get('pairs', {})\n",
    "\n",
    "print(f\"Loaded {len(alias_mappings)} photolysis rate alias mappings\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6a7b8c9",
   "metadata": {},
   "source": [
    "Next, let's create the convenience function. The `tuv-x` altitude grid starts at 0 km. We skip the 0 km level (index 0) because it represents the surface boundary. We use grid edges at indices 1–9, which correspond to 1–10 km."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7b8c9d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_tuv_rates(utc_time, num_grid_cells, start=1):\n",
    "    \"\"\"Return photolysis rate constants for `num_grid_cells` altitude cells starting at index `start`.\"\"\"\n",
    "    lat, lon = boulder\n",
    "    solpos = pvlib.solarposition.get_solarposition(time=utc_time, latitude=lat, longitude=lon)\n",
    "    sza_deg = solpos['zenith'].item()\n",
    "\n",
    "    tuv_rates = tuvx.run(sza=np.deg2rad(sza_deg), earth_sun_distance=1.0)\n",
    "    end = start + num_grid_cells\n",
    "\n",
    "    photolysis_rate_constants = {}\n",
    "    for mapping in alias_mappings:\n",
    "        label    = mapping['to']\n",
    "        scale    = mapping.get(\"scale by\", 1)\n",
    "        tuv_label = mapping['from']\n",
    "        rate = tuv_rates.sel(reaction=tuv_label).photolysis_rate_constants.values * scale\n",
    "        photolysis_rate_constants[f'USER.{label}'] = rate[start:end]\n",
    "\n",
    "    return photolysis_rate_constants, tuv_rates"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8c9d0e1",
   "metadata": {},
   "source": [
    "## 4. Altitude Grid, Atmospheric Conditions, and Solver Setup\n",
    "\n",
    "We use TUV-x altitude grid cells 1–9 (approximately 1–10 km). Temperature and pressure at those altitudes come from the US Standard Atmosphere 1976 model. We compute the initial photolysis rates for the simulation start time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9d0e1f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "start = 1\n",
    "num_grid_cells = 9\n",
    "\n",
    "# Compute initial photolysis rates and grab the altitude edges for our grid cells\n",
    "photolysis_rate_constants, tuv_rates = get_tuv_rates(sim_time, num_grid_cells, start)\n",
    "vertical_edges = tuv_rates.vertical_edge[start:start + num_grid_cells].data\n",
    "\n",
    "print(f\"Altitude range: {vertical_edges.min():.1f} – {vertical_edges.max():.1f} km\")\n",
    "\n",
    "# US Standard Atmosphere T and P at those altitudes (km -> m)\n",
    "atm         = ussa1976.compute(z=vertical_edges * 1000, variables=[\"t\", \"p\"])\n",
    "temperature = atm['t'].values\n",
    "pressure    = atm['p'].values\n",
    "\n",
    "print(f\"Temperature range: {temperature.min():.1f} – {temperature.max():.1f} K\")\n",
    "print(f\"Pressure range:    {pressure.min():.1f} – {pressure.max():.1f} Pa\")\n",
    "\n",
    "# Load mechanism and create solver\n",
    "parser = Parser()\n",
    "mechanism = parser.parse(find_config_path(\"v1\", \"ts1\", \"ts1.json\"))\n",
    "\n",
    "solver = musica.MICM(\n",
    "    mechanism=mechanism,\n",
    "    solver_type=musica.SolverType.rosenbrock_standard_order\n",
    ")\n",
    "\n",
    "state = solver.create_state(num_grid_cells)\n",
    "\n",
    "print(f\"Solver created with {num_grid_cells} grid cells\")\n",
    "print(f\"Mechanism has {len(list(state.get_species_ordering()))} species\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0e1f2a3",
   "metadata": {},
   "source": [
    "## 5. Load Initial Conditions\n",
    "\n",
    "The initial conditions CSV contains parameters for all grid cells. Each row has a parameter name and one or two values, following this convention:\n",
    "\n",
    "| Prefix | value1 | value2 |\n",
    "|--------|--------|--------|\n",
    "| `CONC` | Initial concentration (mol m⁻³) | — |\n",
    "| `ENV`  | Temperature (K) or Pressure (Pa) | — |\n",
    "| `USER` | User-defined parameter value | — |\n",
    "| `SURF` | Effective radius (m) | Particle number concentration (# m⁻³) |\n",
    "\n",
    "We split the file into these four groups and process each separately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1f2a3b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "conditions = pd.read_csv(\n",
    "    find_config_path(\"v1\", \"ts1\", \"initial_conditions.csv\"),\n",
    "    sep=',',\n",
    "    names=['parameter', 'value1', 'value2'],\n",
    "    dtype={'parameter': str, 'value1': float, 'value2': float}\n",
    ")\n",
    "\n",
    "surface_reactions       = conditions[conditions['parameter'].str.contains('SURF')]\n",
    "initial_concentrations  = conditions[conditions['parameter'].str.contains('CONC')].copy()\n",
    "environmental_conditions = conditions[conditions['parameter'].str.contains('ENV')]\n",
    "user_defined_conditions = conditions[conditions['parameter'].str.contains('USER')]\n",
    "\n",
    "# Strip the 'CONC.' prefix so names match the species ordering expected by the solver\n",
    "initial_concentrations.loc[:, 'parameter'] = (\n",
    "    initial_concentrations['parameter'].str.replace('CONC.', '', regex=False)\n",
    ")\n",
    "\n",
    "assert (\n",
    "    len(surface_reactions) + len(initial_concentrations) +\n",
    "    len(environmental_conditions) + len(user_defined_conditions)\n",
    ") == len(conditions)\n",
    "\n",
    "print(f\"Loaded {len(initial_concentrations)} species, {len(user_defined_conditions)} user-defined params, \"\n",
    "      f\"{len(surface_reactions)} surface reactions\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2a3b4c5",
   "metadata": {},
   "source": [
    "## 6. Set State Conditions\n",
    "\n",
    "We populate the solver state with:\n",
    "- **Concentrations** from the initial conditions CSV, broadcast to all grid cells.\n",
    "- **Temperature and pressure** from the US Standard Atmosphere (`ussa1976`), evaluated at the `tuv-x` altitude grid edges for our 9 cells.\n",
    "- **User-defined rate parameters** from the CSV, overwritten with the TUV-x photolysis rates where applicable.\n",
    "- **Surface reaction parameters** (effective radius and particle number concentration) also from the CSV."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3b4c5d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build concentration dict — same value broadcast to every grid cell\n",
    "concentration_dict = {\n",
    "    row['parameter']: [row['value1']] * num_grid_cells\n",
    "    for _, row in initial_concentrations.iterrows()\n",
    "}\n",
    "\n",
    "# Build user-defined rate parameters dict\n",
    "user_defined_dict = {\n",
    "    row['parameter']: [row['value1']] * num_grid_cells\n",
    "    for _, row in user_defined_conditions.iterrows()\n",
    "}\n",
    "\n",
    "# Surface reactions require two parameters: effective radius and particle number concentration\n",
    "for _, row in surface_reactions.iterrows():\n",
    "    user_defined_dict[f\"{row['parameter']}.effective radius [m]\"] = [row['value1']] * num_grid_cells\n",
    "    user_defined_dict[f\"{row['parameter']}.particle number concentration [# m-3]\"] = [row['value2']] * num_grid_cells\n",
    "\n",
    "# Overwrite scalar USER.j* values with the per-altitude TUV-x photolysis rates\n",
    "found_rates  = sorted(photolysis_rate_constants.keys())\n",
    "needed_rates = sorted([k for k in state.get_user_defined_rate_parameters() if 'USER.j' in k])\n",
    "missing_rates = set(needed_rates) - set(found_rates)\n",
    "if missing_rates:\n",
    "    print(f\"Missing photolysis rates (will remain at CSV values): {missing_rates}\")\n",
    "user_defined_dict.update(photolysis_rate_constants)\n",
    "\n",
    "# Apply everything to the state\n",
    "state.set_conditions(temperature, pressure)\n",
    "state.set_concentrations(concentration_dict)\n",
    "state.set_user_defined_rate_parameters(user_defined_dict)\n",
    "\n",
    "print(\"State initialised.\")\n",
    "print(f\"Temperature range: {temperature.min():.1f} – {temperature.max():.1f} K\")\n",
    "print(f\"Pressure range:    {pressure.min():.0f} – {pressure.max():.0f} Pa\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4c5d6e7",
   "metadata": {},
   "source": [
    "## 7. Run the Simulation\n",
    "\n",
    "We integrate forward for 1/10 of a day (~2.4 hours) in 30-second steps. At the end of each step we:\n",
    "- Advance the simulation clock by the step duration.\n",
    "- Recompute the solar zenith angle for the new time using `pvlib`.\n",
    "- Update the photolysis rates on the state before the next step.\n",
    "\n",
    "This means the photolysis rates change with the diurnal cycle throughout the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5d6e7f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_times          = [sim_time]\n",
    "concentrations     = [state.get_concentrations()]\n",
    "photo_rate_history = [{k: v.copy() for k, v in state.get_user_defined_rate_parameters().items()}]\n",
    "\n",
    "time_step         = 30                        # seconds per output step\n",
    "simulation_length = 0.1 * 24 * SECONDS_PER_HOUR  # 1/10 of a day in seconds\n",
    "current_time      = 0\n",
    "last_printed_percent = -5\n",
    "\n",
    "while current_time < simulation_length:\n",
    "    # Inner loop: accumulate solver sub-steps until a full time_step is consumed\n",
    "    elapsed = 0\n",
    "    while elapsed < time_step:\n",
    "        remaining = time_step - elapsed\n",
    "        result = solver.solve(state, remaining)\n",
    "        elapsed      += result.stats.final_time\n",
    "        current_time += result.stats.final_time\n",
    "        if result.state != SolverState.Converged:\n",
    "            print(f\"  Solver state: {result.state} at t={current_time:.1f} s\")\n",
    "\n",
    "    # Advance the wall-clock time and update photolysis rates for the new position\n",
    "    sim_time += timedelta(seconds=time_step)\n",
    "    photolysis_rate_constants, _ = get_tuv_rates(sim_time, num_grid_cells, start)\n",
    "\n",
    "    user_defined_dict = state.get_user_defined_rate_parameters()\n",
    "    for key in photolysis_rate_constants:\n",
    "        user_defined_dict[key] = photolysis_rate_constants[key]\n",
    "    state.set_user_defined_rate_parameters(user_defined_dict)\n",
    "\n",
    "    sim_times.append(sim_time)\n",
    "    concentrations.append(state.get_concentrations())\n",
    "    photo_rate_history.append({k: v.copy() for k, v in state.get_user_defined_rate_parameters().items()})\n",
    "\n",
    "    current_percent = (current_time / simulation_length) * 100\n",
    "    rounded = int(current_percent // 5) * 5\n",
    "    if rounded > last_printed_percent:\n",
    "        last_printed_percent = rounded\n",
    "        local_str = sim_time.astimezone(boulder_tz).strftime('%H:%M %Z')\n",
    "        print(f\"Progress: {last_printed_percent}%  (sim time: {local_str})\")\n",
    "\n",
    "print(\"Simulation complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6e7f8a9",
   "metadata": {},
   "source": [
    "## 8. Collect Results into an xarray Dataset\n",
    "\n",
    "We organise the simulation output into an `xr.Dataset` with dimensions `time` and `height`. Time is stored as actual `datetime64` timestamps. Photolysis rates are recorded at every time step so we can see how they vary with the diurnal cycle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7f8a9b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_conditions = state.get_conditions()\n",
    "\n",
    "data_vars = {\n",
    "    \"temperature\": ([\"height\"], final_conditions[\"temperature\"], {\"units\": \"K\"}),\n",
    "    \"pressure\":    ([\"height\"], final_conditions[\"pressure\"],    {\"units\": \"Pa\"}),\n",
    "    \"air_density\": ([\"height\"], final_conditions[\"air_density\"], {\"units\": \"mol m-3\"}),\n",
    "}\n",
    "\n",
    "# Photolysis rates over time\n",
    "photo_names  = list(photo_rate_history[0].keys())\n",
    "photo_values = [[params[name] for name in photo_names] for params in photo_rate_history]\n",
    "data_vars[\"photolysis_rates\"] = (\n",
    "    [\"time\", \"photolysis_rate_names\", \"height\"],\n",
    "    photo_values,\n",
    "    {\"units\": \"s-1\"}\n",
    ")\n",
    "\n",
    "# Species concentrations (time x height)\n",
    "species_ordering = state.get_species_ordering()\n",
    "for species in species_ordering:\n",
    "    data_vars[species] = (\n",
    "        [\"time\", \"height\"],\n",
    "        np.array([c[species] for c in concentrations], dtype=np.float64),\n",
    "        {\"units\": \"mol m-3\"}\n",
    "    )\n",
    "\n",
    "ds = xr.Dataset(\n",
    "    data_vars,\n",
    "    coords={\n",
    "        \"time\":                  np.array(sim_times, dtype=\"datetime64[s]\"),\n",
    "        \"height\":                vertical_edges,\n",
    "        \"photolysis_rate_names\": np.array(photo_names, dtype=\"S\"),\n",
    "    }\n",
    ")\n",
    "\n",
    "print(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8a9b0c1",
   "metadata": {},
   "source": [
    "## 9. Plot Results\n",
    "\n",
    "We plot the time evolution of four representative species — BEPOMUC, C6H5OOH, BR, and CL — for every grid cell. Each line corresponds to one altitude."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9b0c1d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "elapsed_hours = (ds['time'] - ds['time'].isel(time=0)) / np.timedelta64(1, 'h')\n",
    "species_to_plot = ['BEPOMUC', 'C6H5OOH', 'BR', 'CL']\n",
    "\n",
    "fig, axes = plt.subplots(2, 2, figsize=(12, 8))\n",
    "\n",
    "for ax, species in zip(axes.flat, species_to_plot):\n",
    "    for cell in range(num_grid_cells):\n",
    "        ax.plot(elapsed_hours, ds[species].isel(height=cell), label=f\"{vertical_edges[cell]:.1f} km\")\n",
    "    ax.set_title(species)\n",
    "    ax.set_xlabel('Time [hours since start]')\n",
    "    ax.set_ylabel('Concentration [mol m⁻³]')\n",
    "    ax.set_xlim(0, simulation_length / SECONDS_PER_HOUR)\n",
    "    ax.set_ylim(0, None)\n",
    "    ax.grid(True, alpha=0.5)\n",
    "    ax.spines[:].set_visible(False)\n",
    "    ax.tick_params(width=0)\n",
    "\n",
    "# Add a shared legend outside the last axis\n",
    "handles, labels = axes.flat[-1].get_legend_handles_labels()\n",
    "fig.legend(handles, labels, title='Altitude', loc='lower center', ncol=num_grid_cells, bbox_to_anchor=(0.5, -0.05))\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.12.3)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}